import requests
import json
import time
from typing import Dict, Any, Optional

class LLMAgent:
    def __init__(
        self,
        model_type: str = "local",  # "local" 或 "online"
        base_url: str = "http://localhost:1234/v1",
        api_key: Optional[str] = None,
        online_provider: str = "openai",  # 目前支持 "openai"
        default_model: Optional[str] = None
    ):
        """
        初始化LLM Agent
        :param model_type: 模型类型，"local" 表示本地LM Studio模型，"online" 表示在线大模型
        :param base_url: LM Studio服务的基础URL，仅当model_type="local"时有效
        :param api_key: 在线大模型的API密钥，仅当model_type="online"时有效
        :param online_provider: 在线大模型提供商，目前仅支持"openai"
        :param default_model: 默认使用的模型名称
        """
        self.model_type = model_type.lower()
        self.base_url = base_url
        self.api_key = api_key
        self.online_provider = online_provider.lower()
        self.default_model = default_model
        self.headers = {
            "Content-Type": "application/json"
        }

        # 验证参数
        if self.model_type == "online":
            if not self.api_key:
                raise ValueError("使用在线大模型时必须提供API密钥(api_key)")
            if self.online_provider == "openai":
                self.headers["Authorization"] = f"Bearer {self.api_key}"
                self.base_url = "https://api.openai.com/v1"
            else:
                raise NotImplementedError(f"暂不支持的在线大模型提供商: {self.online_provider}")
        elif self.model_type != "local":
            raise ValueError(f"不支持的模型类型: {self.model_type}，仅支持 'local' 和 'online'")

        # 检查服务是否可用
        if self.model_type == "local":
            self._check_service_availability()
        else:
            print(f"✅ 已配置为使用{self.online_provider}在线大模型")

    def _check_service_availability(self):
        """检查LM Studio服务是否可用"""
        try:
            response = requests.get(f"{self.base_url}/models")
            if response.status_code != 200:
                raise ConnectionError(f"LM Studio服务不可用，状态码: {response.status_code}")
            print("✅ LM Studio服务已连接成功")
        except Exception as e:
            print(f"⚠️ 警告: 无法连接到LM Studio服务: {e}")
            print(f"请确保LM Studio已启动并运行在 {self.base_url}")

    def list_models(self) -> Dict[str, Any]:
        """列出可用的模型"""
        try:
            if self.model_type == "online" and self.online_provider == "openai":
                # OpenAI API 获取模型列表的端点
                response = requests.get(f"{self.base_url}/models", headers=self.headers)
            else:
                response = requests.get(f"{self.base_url}/models")
            response.raise_for_status()
            return response.json()
        except Exception as e:
            print(f"获取模型列表失败: {e}")
            return {"error": str(e)}

    def chat_completion(
        self,
        messages: list,
        model: Optional[str] = None,
        temperature: float = 0.7,
        max_tokens: int = 1000,
        stream: bool = False
    ) -> Dict[str, Any]:
        """
        发送聊天完成请求
        :param messages: 消息列表，格式为 [{"role": "user", "content": "你好"}]
        :param model: 模型名称，如果为None则使用默认模型
        :param temperature: 温度参数，控制输出的随机性
        :param max_tokens: 最大生成 tokens 数
        :param stream: 是否流式返回结果
        :return: 响应结果
        """
        payload = {
            "messages": messages,
            "temperature": temperature,
            "max_tokens": max_tokens,
            "stream": stream
        }

        # 使用默认模型（如果提供）
        if model:
            payload["model"] = model
        elif self.default_model:
            payload["model"] = self.default_model

        try:
            if self.model_type == "online" and self.online_provider == "openai":
                response = requests.post(
                    f"{self.base_url}/chat/completions",
                    headers=self.headers,
                    data=json.dumps(payload)
                )
            else:
                response = requests.post(
                    f"{self.base_url}/chat/completions",
                    headers=self.headers,
                    data=json.dumps(payload)
                )
            response.raise_for_status()
            return response.json()
        except Exception as e:
            print(f"聊天完成请求失败: {e}")
            return {"error": str(e)}

    def generate_text(
        self,
        prompt: str,
        model: Optional[str] = None,
        temperature: float = 0.7,
        max_tokens: int = 1000
    ) -> str:
        """
        生成文本
        :param prompt: 提示文本
        :param model: 模型名称，如果为None则使用默认模型
        :param temperature: 温度参数
        :param max_tokens: 最大生成 tokens 数
        :return: 生成的文本
        """
        messages = [{"role": "user", "content": prompt}]
        result = self.chat_completion(messages, model, temperature, max_tokens)

        if "error" in result:
            return f"生成文本失败: {result['error']}"

        if "choices" in result and len(result["choices"]) > 0:
            return result["choices"][0]["message"]["content"]
        else:
            return "没有生成任何文本"

# 示例用法
if __name__ == "__main__":
    # 使用本地LM Studio模型
    print("\n===== 使用本地LM Studio模型 ======")
    try:
        local_agent = LLMAgent(model_type="local")
        local_models = local_agent.list_models()
        print("本地模型列表:", local_models)
        local_response = local_agent.generate_text("请解释什么是人工智能")
        print(f"本地模型响应: {local_response}")
    except Exception as e:
        print(f"本地模型示例失败: {e}")

    # 使用在线大模型（需要API密钥）
    print("\n===== 使用在线大模型 ======")
    try:
        # 注意：这里需要替换为实际的API密钥
        api_key = "your_api_key_here"
        online_agent = LLMAgent(
            model_type="online",
            api_key=api_key,
            online_provider="openai",
            default_model="gpt-3.5-turbo"
        )
        online_models = online_agent.list_models()
        print("在线模型列表: (仅显示前5个)", online_models["data"][:5] if "data" in online_models else online_models)
        online_response = online_agent.generate_text("请解释什么是人工智能")
        print(f"在线模型响应: {online_response}")
    except Exception as e:
        print(f"在线模型示例失败: {e}")
        print("提示: 请替换api_key为有效的API密钥，并确保网络连接正常。")